package org.tinygroup.daccase;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.commons.lang.math.RandomUtils;
import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.protocol.java.sampler.AbstractJavaSamplerClient;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;
import org.apache.jmeter.samplers.SampleResult;
import org.tinygroup.dac.context.DataAccessControlContext;
import org.tinygroup.dac.context.impl.DefaultDataAccessControlContext;

public class DACTest extends AbstractJavaSamplerClient {

	private static final int DEFAULT_BATCH_SIZE = 1000;
	private static DACTestInitial init;
	static {
		init = DACTestInitial.getInstance();
		init.init();
	}

	private String isDac;
	private String testMode;
	private int batchNum;

	private static volatile int executeTimes = 0;
	private static AtomicInteger atomicInteger = new AtomicInteger(0);

	@Override
	public void setupTest(JavaSamplerContext context) {
		super.setupTest(context);
		this.isDac = context.getParameter("isDac", "true");// 是否分片
		this.testMode = context.getParameter("testMode", "dataPrepare");// 测试方法
		this.batchNum = context.getIntParameter("batchNum");// 测试方法
		if (testMode.equals("dataPrepare")) {
			Connection connection = null;
			try {
				if (isDac.equals("true")) {
					connection = init.getDacConnection();
				} else {
					connection = init.getCommonConnection();
				}
				deleteOperate("delete from t_user", connection);
				deleteOperate("delete from t_function", connection);
				deleteOperate("delete from t_user_function", connection);
			} catch (SQLException e) {
				throw new RuntimeException(e);
			} finally {
				try {
					if(connection!=null){
						connection.close();
					}
				} catch (SQLException e) {
				}
			}
		}

	}

	@Override
	public void teardownTest(JavaSamplerContext context) {
		super.teardownTest(context);
		try {
			init.getCommonDatasource().close();
		} catch (SQLException e) {
			e.printStackTrace();
		}
	}

	public SampleResult runTest(JavaSamplerContext context) {
		SampleResult sr = new SampleResult();
		sr.setSampleLabel("dac请求");
		Connection connection = null;
		try {
			if (isDac.equals("true")) {
				connection = init.getDacConnection();
			} else {
				connection = init.getCommonConnection();
			}
			sr.sampleStart();// jmeter 开始统计响应时间标记
			if (testMode.equals("dataPrepare")) {
				dataPrepare(connection);
			} else if (testMode.equals("insert")) {
				insertTest(connection);
			} else if (testMode.equals("update")) {
				updateTest(connection);
			} else if (testMode.equals("delete")) {
				deleteTest(connection);
			} else if (testMode.equals("selectByKey")) {
				selectByKeyTest(connection);
			} else if (testMode.equals("selectFunction")) {
				selectFunctionTest(connection);
			} else if (testMode.equals("selectOrdeyBy")) {
				selectOrderByTest(connection);
			} else if (testMode.equals("selectGroupBy")) {
				selectGroupByTest(connection);
			}
			sr.sampleEnd();// jmeter 结束统计响应时间标记
			sr.setSuccessful(true);
		} catch (SQLException e) {
			sr.setSuccessful(false);
			e.printStackTrace();
		} finally {
			if (connection != null) {
				try {
					connection.close();
				} catch (SQLException e) {
				}
			}
		}
		return sr;
	}

	/**
	 * 数据准备
	 * 
	 * @param connection
	 * @throws SQLException
	 */
	private void dataPrepare(Connection connection) throws SQLException {
		try {
			connection.setAutoCommit(false);
			PreparedStatement preparedStatement = connection
					.prepareStatement("insert into t_user(user_id,user_name,login_name,age,password,salary) values(?,?,?,?,?,?)");
			for (int i = 0; i < batchNum; i++) {
				preparedStatement.setInt(1, i);
				preparedStatement.setString(2, "username" + i);
				preparedStatement.setString(3, "loginName" + i);
				preparedStatement.setInt(4, RandomUtils.nextInt(100));
				preparedStatement.setString(5, "88888888");
				preparedStatement.setInt(6, 10000 + i);
				preparedStatement.addBatch();
				if (i > 0 && i % DEFAULT_BATCH_SIZE == 0) {
					preparedStatement.executeBatch();
				}
			}
			preparedStatement.executeBatch();
			preparedStatement.close();
			preparedStatement = connection
					.prepareStatement("insert into t_function(function_id,function_name,function_code) values(?,?,?)");
			for (int i = 0; i < batchNum; i++) {
				preparedStatement.setInt(1, 1000 + i);
				preparedStatement.setString(2, "function_name" + i);
				preparedStatement.setString(3, "function_code" + i);
				preparedStatement.addBatch();
				if (i > 0 && i % DEFAULT_BATCH_SIZE == 0) {
					preparedStatement.executeBatch();
				}
			}
			preparedStatement.executeBatch();
			preparedStatement.close();
			preparedStatement = connection
					.prepareStatement("insert into t_user_function(id,user_id,function_id) values(?,?,?)");
			int primaryValue = 0;
			for (int i = 0; i < batchNum; i++) {
				for (int j = 0; j < 10; j++) {
					preparedStatement.setInt(1, ++primaryValue);
					preparedStatement.setInt(2, i);
					preparedStatement.setInt(3, i + 1000 + j);
					preparedStatement.addBatch();
					if (primaryValue % DEFAULT_BATCH_SIZE == 0) {
						preparedStatement.executeBatch();
					}
				}
			}
			preparedStatement.executeBatch();
			preparedStatement.close();
			connection.commit();
		} catch (SQLException e) {
			connection.rollback();
			throw e;
		}
	}

	private void deleteTest(Connection connection) throws SQLException {
		PreparedStatement preparedStatement = connection
				.prepareStatement("delete from t_user where user_id=?");
		try {
			int next = executeTimes++;
			// setUserNameParms(next);
			preparedStatement.setInt(1, next);
			preparedStatement.executeUpdate();
		} finally {
			if (preparedStatement != null) {
				preparedStatement.close();
			}
		}

	}

	private void insertTest(Connection connection) throws SQLException {
		PreparedStatement preparedStatement = connection
				.prepareStatement("insert into t_user(user_id,user_name,login_name,age,password,salary) values(?,?,?,?,?,?)");
		try {
			int next = batchNum + atomicInteger.getAndIncrement();
			preparedStatement.setInt(1, next);
			preparedStatement.setString(2, "user_name" + next);
			preparedStatement.setString(3, "loginName" + next);
			preparedStatement.setInt(4, RandomUtils.nextInt(100));
			preparedStatement.setString(5, "888888");
			preparedStatement.setInt(6, 10000 + next);
			preparedStatement.executeUpdate();
		} finally {
			if (preparedStatement != null) {
				preparedStatement.close();
			}
		}

	}

	private void updateTest(Connection connection) throws SQLException {
		PreparedStatement preparedStatement = connection
				.prepareStatement("update t_user set password=? where user_id=?");
		try {
			preparedStatement.setString(1, "xxxxxxxx");
			int next = executeTimes++;
			preparedStatement.setInt(2, next);
			// setUserNameParms(next);
			preparedStatement.executeUpdate();
		} finally {
			if (preparedStatement != null) {
				preparedStatement.close();
			}
		}

	}

	private void selectByKeyTest(Connection connection) throws SQLException {
		PreparedStatement preparedStatement = connection
				.prepareStatement("select * from t_user  where user_id=?");
		try {
			int next = executeTimes++;
			// setUserNameParms(next);
			preparedStatement.setInt(1, next);
			ResultSet rs = preparedStatement.executeQuery();
			if (rs.next()) {
				System.out.println("selectByKeyTest:user_id:"
						+ rs.getLong("user_id"));
			}
			rs.close();
		} finally {
			if (preparedStatement != null) {
				preparedStatement.close();
			}
		}
	}

	private void setUserNameParms(int next) {
		DataAccessControlContext context = DefaultDataAccessControlContext
				.getInstance();
		context.setThreadContextValue("user_name", "user_name" + next);
	}

	private void selectFunctionTest(Connection connection) throws SQLException {
		PreparedStatement preparedStatement = connection
				.prepareStatement("select * from t_function where function_id=?");
		try {
			int next = executeTimes++;
//			DataAccessControlContext context = DefaultDataAccessControlContext
//					.getInstance();
//			context.setThreadContextValue("user_id", next);
			preparedStatement.setInt(1, 100+next);
			ResultSet rs = preparedStatement.executeQuery();
			if (rs.next()) {
				System.out.println("selectFunctionTest:function_id:"
						+ rs.getLong("function_id"));
			}
			rs.close();
		} finally {
			if (preparedStatement != null) {
				preparedStatement.close();
			}
		}
	}

	private void selectOrderByTest(Connection connection) throws SQLException {
		PreparedStatement preparedStatement = connection
				.prepareStatement("select * from t_user  where user_id=? order by user_id");
		try {
			int next = executeTimes++;
			// setUserNameParms(next);
			preparedStatement.setInt(1, next);
			ResultSet rs = preparedStatement.executeQuery();
			if (rs.next()) {
				System.out.println("selectOrderByTest:user_id:"
						+ rs.getLong("user_id"));
			}
			rs.close();
		} finally {
			if (preparedStatement != null) {
				preparedStatement.close();
			}
		}
	}

	/**
	 * 1个线程执行
	 * 
	 * @param connection
	 * @throws SQLException
	 */
	private void selectGroupByTest(Connection connection) throws SQLException {
		PreparedStatement preparedStatement = connection
				.prepareStatement("SELECT SUM(salary) AS salary_sum, user_id FROM t_user GROUP BY user_id ORDER BY salary_sum DESC");
		try {
			ResultSet rs = preparedStatement.executeQuery();
			if (rs.next()) {
				System.out.println("selectGroupByTest:salary_sum:"
						+ rs.getLong("salary_sum"));
			}
			rs.close();
		} finally {
			if (preparedStatement != null) {
				preparedStatement.close();
			}
		}
	}

	private void deleteOperate(String sql, Connection connection)
			throws SQLException {
		Statement statement = connection.createStatement();
		statement.execute(sql);
		statement.close();
	}

	public Arguments getDefaultParameters() {
		Arguments params = new Arguments();
		params.addArgument("isDac", "true");
		params.addArgument("testMode", "dataPrepare");
		params.addArgument("batchNum", "10000");
		return params;
	}

	public static void main(String[] args) { // TODO Auto-generated method stub
		Arguments params = new Arguments();
		params.addArgument("isDac", "true");// 设置参数，并赋予默认值1
		params.addArgument("testMode", "selectByKey");// 设置参数，并赋予默认值1
		params.addArgument("batchNum", "1000");// 设置参数，并赋予默认值2
		JavaSamplerContext arg0 = new JavaSamplerContext(params);
		DACTest test = new DACTest();
		test.setupTest(arg0);
		long startTime = System.currentTimeMillis();
		test.runTest(arg0);
		System.out.println(System.currentTimeMillis() - startTime);
		test.teardownTest(arg0);

	}

}
